from PIL import Image
import numpy as np

FILE_NAME = r'pic.png'
IMG_SIZE = (32, 32)


class ImageError(Exception):
    __metaclass__ = object

    def __init__(self, *args, **kw):
        text = 'The image seems not like a digit.'
        super(ImageError, self).__init__(text, *args, **kw)


def get_box(img_arr):
    """Get the box of the digit. Return the left, top, right, bottom"""
    shape = img_arr.shape
    for i in range(shape[0]):
        if img_arr[i].max() == 255:
            top = i
            break

    for i in range(shape[0]-1, -1, -1):
        if img_arr[i].max() == 255:
            bottom = i
            break

    for i in range(shape[1]):
        if img_arr[:, i].max() == 255:
            left = i
            break

    for i in range(shape[1]-1, -1, -1):
        if img_arr[:, i].max() == 255:
            right = i
            break
    try:
        box = left, top, right, bottom
    except NameError as e:
        print("This shouldn't have happened.")
        raise ImageError

    return box


def pre_process(file_name):
    """Resize the image, move the digit to the center of image according its box.
    Finally return a PIL.Image object of 32x32 size."""
    im = Image.open(file_name)
    im = im.split()[3]  # only the A band takes the information
    # im2 = im.resize(IMG_SIZE)  # to fasten following procedures, assume got a square image

    im_arr = np.asarray(im)
    left, top, right, bottom = get_box(im_arr)

    # stretch the image's height to standard IMG_SIZE
    size = (right - left, bottom - top)
    ratio = IMG_SIZE[1] / float(size[1])
    new_size = int(ratio * size[0]), IMG_SIZE[1]
    center = (new_size[0]/2, new_size[1]/2)
    rel = int(IMG_SIZE[0]/2 - center[0])  # x relative move of the center of IMG_SIZE and new_size

    print("center:%s ratio:%f size:%s new_size:%s" %
          (str(center), ratio, str(size), str(new_size)))
    im = im.crop((left, top, right, bottom))
    im = im.resize(new_size)
    
    out = Image.new('L', IMG_SIZE, 0)
    out.paste(im, (rel, 0))

    '''
    im = im.crop((left, top, right, bottom))
    out = im.resize(IMG_SIZE)
    '''

    return out


def get_array(img):
    """Convert the image to the numpy array."""
    arr = np.asarray(img).reshape([1, 1024])
    arr = arr/255

    return arr


if __name__ == '__main__':
    import knn_recog

    im = pre_process(FILE_NAME)
    im_arr = get_array(im)
    print(knn_recog.predict(im_arr))




